"""
MCP Client 实现
与MCP服务器交互的客户端
"""
import asyncio
import json
from datetime import datetime
from typing import Any, Dict, List, Optional

import httpx
import websockets
from websockets.exceptions import ConnectionClosedError

from mcp_models import (
    MCPCapabilities, ToolCall, ToolResult, 
    ResourceRequest, ResourceResponse, PromptRequest, PromptResponse
)


class MCPClient:
    """MCP客户端实现"""
    
    def __init__(self, server_url: str = "http://localhost:8000"):
        self.server_url = server_url
        self.ws_url = server_url.replace("http://", "ws://").replace("https://", "wss://") + "/ws"
        self.capabilities: Optional[MCPCapabilities] = None
        self.websocket: Optional[websockets.WebSocketServerProtocol] = None
        
    async def connect(self):
        """连接到MCP服务器"""
        try:
            # 获取服务器能力
            await self.get_capabilities()
            
            # 建立WebSocket连接
            await self.connect_websocket()
            
            print(f"✅ Successfully connected to MCP Server at {self.server_url}")
            return True
        except Exception as e:
            print(f"❌ Failed to connect to server: {e}")
            return False
    
    async def get_capabilities(self) -> MCPCapabilities:
        """获取服务器能力"""
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{self.server_url}/capabilities")
            response.raise_for_status()
            self.capabilities = MCPCapabilities(**response.json())
            return self.capabilities
    
    async def connect_websocket(self):
        """建立WebSocket连接"""
        try:
            self.websocket = await websockets.connect(self.ws_url)
            print("🔗 WebSocket connected")
        except Exception as e:
            print(f"⚠️ WebSocket connection failed: {e}")
    
    async def disconnect(self):
        """断开连接"""
        if self.websocket:
            await self.websocket.close()
            print("🔌 WebSocket disconnected")
    
    async def call_tool(self, tool_name: str, parameters: Dict[str, Any], request_id: Optional[str] = None) -> ToolResult:
        """调用工具"""
        tool_call = ToolCall(
            tool_name=tool_name,
            parameters=parameters,
            request_id=request_id or f"req_{datetime.now().timestamp()}"
        )
        
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.server_url}/tools/call",
                json=tool_call.model_dump()
            )
            response.raise_for_status()
            return ToolResult(**response.json())
    
    async def get_resource(self, uri: str, parameters: Optional[Dict[str, Any]] = None) -> ResourceResponse:
        """获取资源"""
        resource_request = ResourceRequest(uri=uri, parameters=parameters)
        
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.server_url}/resources/get",
                json=resource_request.model_dump()
            )
            response.raise_for_status()
            return ResourceResponse(**response.json())
    
    async def render_prompt(self, name: str, parameters: Dict[str, Any]) -> PromptResponse:
        """渲染提示"""
        prompt_request = PromptRequest(name=name, parameters=parameters)
        
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.server_url}/prompts/render",
                json=prompt_request.model_dump()
            )
            response.raise_for_status()
            return PromptResponse(**response.json())
    
    async def listen_notifications(self):
        """监听WebSocket通知"""
        if not self.websocket:
            print("⚠️ WebSocket not connected")
            return
        
        try:
            async for message in self.websocket:
                data = json.loads(message)
                await self.handle_websocket_message(data)
        except ConnectionClosedError:
            print("🔌 WebSocket connection closed")
        except Exception as e:
            print(f"❌ WebSocket error: {e}")
    
    async def handle_websocket_message(self, data: Dict[str, Any]):
        """处理WebSocket消息"""
        msg_type = data.get("type")
        
        if msg_type == "welcome":
            print(f"🎉 {data.get('message')}")
            capabilities = data.get("capabilities", {})
            print(f"📋 Server capabilities: {capabilities}")
        
        elif msg_type == "notification":
            notification = data.get("data", {})
            priority = notification.get("priority", "normal")
            message = notification.get("message", "")
            
            priority_emoji = {"high": "🔴", "normal": "🔵", "low": "⚪"}.get(priority, "🔵")
            print(f"{priority_emoji} Notification: {message}")
        
        elif msg_type == "pong":
            print("🏓 Pong received")
    
    async def send_ping(self):
        """发送ping消息"""
        if self.websocket:
            await self.websocket.send(json.dumps({"type": "ping"}))
    
    def print_capabilities(self):
        """打印服务器能力"""
        if not self.capabilities:
            print("❓ No capabilities loaded")
            return
        
        print("\n📊 MCP Server Capabilities:")
        print("=" * 50)
        
        # 工具
        print(f"\n🔧 Tools ({len(self.capabilities.tools)}):")
        for tool in self.capabilities.tools:
            print(f"  • {tool.name}: {tool.description}")
            if tool.parameters:
                for param in tool.parameters:
                    required = "✓" if param.required else "○"
                    default = f" (default: {param.default})" if param.default else ""
                    print(f"    {required} {param.name} ({param.type}): {param.description}{default}")
        
        # 资源
        print(f"\n📂 Resources ({len(self.capabilities.resources)}):")
        for resource in self.capabilities.resources:
            print(f"  • {resource.name} ({resource.type.value}): {resource.description}")
            print(f"    URI: {resource.uri}")
            if resource.template_params:
                print(f"    Parameters: {resource.template_params}")
        
        # 提示
        print(f"\n💬 Prompts ({len(self.capabilities.prompts)}):")
        for prompt in self.capabilities.prompts:
            print(f"  • {prompt.name}: {prompt.description}")
            if prompt.parameters:
                for param in prompt.parameters:
                    required = "✓" if param.required else "○"
                    default = f" (default: {param.default})" if param.default else ""
                    print(f"    {required} {param.name} ({param.type}): {param.description}{default}")


class MCPDemo:
    """MCP演示类"""
    
    def __init__(self):
        self.client = MCPClient()
    
    async def run_demo(self):
        """运行完整演示"""
        print("🚀 MCP (Model Context Protocol) 演示")
        print("=" * 60)
        
        # 连接服务器
        if not await self.client.connect():
            return
        
        # 显示能力
        self.client.print_capabilities()
        
        # 创建监听任务
        listen_task = asyncio.create_task(self.client.listen_notifications())
        
        try:
            # 演示三个核心基元
            await self.demo_tools()
            await self.demo_resources()
            await self.demo_prompts()
            
            # 综合演示
            await self.demo_integration()
            
            print("\n⏳ 等待3秒来接收通知...")
            await asyncio.sleep(3)
            
        finally:
            listen_task.cancel()
            await self.client.disconnect()
    
    async def demo_tools(self):
        """演示工具功能"""
        print("\n" + "=" * 30)
        print("🔧 工具演示 (Tools - 模型控制)")
        print("=" * 30)
        
        # 1. 发送通知
        print("\n1️⃣ 发送通知工具:")
        result = await self.client.call_tool(
            "send_notification",
            {"message": "MCP演示：工具调用测试", "priority": "high"}
        )
        if result.success:
            print(f"✅ 通知发送成功: {result.result}")
        else:
            print(f"❌ 通知发送失败: {result.error}")
        
        # 2. 创建文件
        print("\n2️⃣ 创建文件工具:")
        file_content = """# MCP 演示文件
这是通过MCP工具创建的文件。

## MCP三个核心基元:
1. 工具 (Tools) - 模型控制
2. 资源 (Resources) - 应用控制  
3. 提示 (Prompts) - 用户控制
"""
        
        result = await self.client.call_tool(
            "create_file",
            {"filename": "mcp_demo.md", "content": file_content}
        )
        if result.success:
            print(f"✅ 文件创建成功: {result.result}")
        else:
            print(f"❌ 文件创建失败: {result.error}")
        
        # 3. 执行命令
        print("\n3️⃣ 执行命令工具:")
        result = await self.client.call_tool(
            "execute_command",
            {"command": "echo 'Hello from MCP Tool!'"}
        )
        if result.success:
            print(f"✅ 命令执行成功:")
            print(f"   输出: {result.result.get('stdout', '').strip()}")
        else:
            print(f"❌ 命令执行失败: {result.error}")
    
    async def demo_resources(self):
        """演示资源功能"""
        print("\n" + "=" * 30)
        print("📂 资源演示 (Resources - 应用控制)")
        print("=" * 30)
        
        # 1. 获取日志文件
        print("\n1️⃣ 获取系统日志资源:")
        today = datetime.now().strftime("%Y-%m-%d")
        try:
            resource = await self.client.get_resource(
                "file://logs/{date}.log",
                {"date": today}
            )
            print(f"✅ 日志获取成功:")
            print(f"   URI: {resource.uri}")
            print(f"   内容类型: {resource.content_type}")
            print(f"   内容预览: {resource.content[:100]}...")
        except Exception as e:
            print(f"❌ 日志获取失败: {e}")
        
        # 2. 获取天气信息
        print("\n2️⃣ 获取天气信息资源:")
        try:
            resource = await self.client.get_resource("api://weather/北京")
            print(f"✅ 天气信息获取成功:")
            print(f"   URI: {resource.uri}")
            print(f"   内容类型: {resource.content_type}")
            weather_data = json.loads(resource.content)
            print(f"   天气信息: {weather_data['city']} - {weather_data['weather']} - {weather_data['temperature']}°C")
        except Exception as e:
            print(f"❌ 天气信息获取失败: {e}")
    
    async def demo_prompts(self):
        """演示提示功能"""
        print("\n" + "=" * 30)
        print("💬 提示演示 (Prompts - 用户控制)")
        print("=" * 30)
        
        # 1. 代码审查提示
        print("\n1️⃣ 代码审查提示:")
        sample_code = """def calculate_fibonacci(n):
    if n <= 1:
        return n
    return calculate_fibonacci(n-1) + calculate_fibonacci(n-2)"""
        
        try:
            prompt_response = await self.client.render_prompt(
                "code_review",
                {
                    "code": sample_code,
                    "language": "python",
                    "focus_areas": "性能优化、算法效率"
                }
            )
            print(f"✅ 代码审查提示生成成功:")
            print(f"{prompt_response.rendered_template}")
        except Exception as e:
            print(f"❌ 提示生成失败: {e}")
        
        # 2. 文档生成提示
        print("\n2️⃣ 文档生成提示:")
        try:
            prompt_response = await self.client.render_prompt(
                "documentation",
                {
                    "content": "async def send_notification(message: str, priority: str = 'normal') -> Dict[str, Any]",
                    "type": "异步函数"
                }
            )
            print(f"✅ 文档生成提示成功:")
            print(f"{prompt_response.rendered_template}")
        except Exception as e:
            print(f"❌ 文档生成失败: {e}")
    
    async def demo_integration(self):
        """综合演示"""
        print("\n" + "=" * 40)
        print("🎯 综合演示 - MCP三个基元协同工作")
        print("=" * 40)
        
        print("\n📝 场景: AI代码助手工作流")
        print("1. 用户通过提示模板请求代码审查")
        print("2. AI使用工具创建审查报告文件")
        print("3. AI通过资源获取相关日志信息")
        print("4. AI发送完成通知")
        
        # 步骤1: 使用提示模板
        code_to_review = """
class DatabaseConnection:
    def __init__(self, host, port, username, password):
        self.host = host
        self.port = port
        self.username = username
        self.password = password  # 明文存储密码！
        
    def connect(self):
        # 没有异常处理
        return f"mysql://{self.username}:{self.password}@{self.host}:{self.port}"
"""
        
        prompt_result = await self.client.render_prompt(
            "code_review",
            {
                "code": code_to_review,
                "language": "python",
                "focus_areas": "安全性漏洞、异常处理"
            }
        )
        
        # 步骤2: 使用工具创建审查报告
        report_content = f"""# 代码审查报告
生成时间: {datetime.now().isoformat()}

## 审查请求
{prompt_result.rendered_template}

## 审查结果
发现以下问题:
1. 🔴 安全风险: 密码明文存储
2. 🟡 异常处理: 缺少连接异常处理
3. 🟡 代码质量: 建议使用环境变量管理配置

## 建议
- 使用加密存储敏感信息
- 添加try-catch异常处理
- 实现连接池管理
"""
        
        file_result = await self.client.call_tool(
            "create_file",
            {
                "filename": f"code_review_report_{int(datetime.now().timestamp())}.md",
                "content": report_content
            }
        )
        
        # 步骤3: 获取相关资源信息
        today = datetime.now().strftime("%Y-%m-%d")
        log_result = await self.client.get_resource(
            "file://logs/{date}.log",
            {"date": today}
        )
        
        # 步骤4: 发送完成通知
        notification_result = await self.client.call_tool(
            "send_notification",
            {
                "message": f"代码审查完成！报告已生成: {file_result.result['filename']}",
                "priority": "normal"
            }
        )
        
        print(f"\n✅ 综合演示完成:")
        print(f"   📋 提示渲染: 成功")
        print(f"   📄 报告生成: {file_result.result['filename']}")
        print(f"   📂 日志获取: {len(log_result.content)} 字符")
        print(f"   🔔 通知发送: 成功")


async def main():
    """主函数"""
    demo = MCPDemo()
    await demo.run_demo()


if __name__ == "__main__":
    asyncio.run(main()) 